package com.cyx.exercise.ollama;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.ExecutorService;

/**
 * Ollama API客户端 - 使用传统URLConnection实现
 * 兼容Java 8+，不依赖java.net.http包
 */
public class OllamaClientLegacy {
    
    private final String baseUrl;
    private final ExecutorService executorService;
    private final int connectTimeout;
    private final int readTimeout;
    
    public OllamaClientLegacy(String baseUrl) {
        this(baseUrl, 30000, 300000); // 30秒连接超时，5分钟读取超时
    }
    
    public OllamaClientLegacy(String baseUrl, int connectTimeout, int readTimeout) {
        this.baseUrl = baseUrl.endsWith("/") ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl;
        this.connectTimeout = connectTimeout;
        this.readTimeout = readTimeout;
        this.executorService = Executors.newCachedThreadPool();
    }
    
    /**
     * 同步调用Ollama生成接口
     */
    public OllamaResponse generate(OllamaRequest request) throws IOException {
        String jsonBody = buildJsonRequest(request);
        String response = sendPostRequest(baseUrl + "/api/generate", jsonBody);
        return parseResponse(response);
    }
    
    /**
     * 异步调用Ollama生成接口
     */
    public CompletableFuture<OllamaResponse> generateAsync(OllamaRequest request) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                return generate(request);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }, executorService);
    }
    
    /**
     * 调用Ollama聊天接口
     */
    public OllamaChatResponse chat(OllamaChatRequest request) throws IOException {
        String jsonBody = buildChatJsonRequest(request);
        String response = sendPostRequest(baseUrl + "/api/chat", jsonBody);
        return parseChatResponse(response);
    }
    
    /**
     * 获取可用模型列表
     */
    public String getModels() throws IOException {
        return sendGetRequest(baseUrl + "/api/tags");
    }
    
    /**
     * 发送POST请求
     */
    private String sendPostRequest(String urlString, String jsonBody) throws IOException {
        URL url = new URL(urlString);
        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
        
        try {
            // 设置请求方法和属性
            connection.setRequestMethod("POST");
            connection.setRequestProperty("Content-Type", "application/json");
            connection.setRequestProperty("Accept", "application/json");
            connection.setConnectTimeout(connectTimeout);
            connection.setReadTimeout(readTimeout);
            connection.setDoOutput(true);
            
            // 发送请求体
            try (OutputStream os = connection.getOutputStream()) {
                byte[] input = jsonBody.getBytes(StandardCharsets.UTF_8);
                os.write(input, 0, input.length);
            }
            
            // 检查响应状态码
            int responseCode = connection.getResponseCode();
            if (responseCode != HttpURLConnection.HTTP_OK) {
                String errorResponse = readErrorStream(connection);
                throw new IOException("HTTP Error: " + responseCode + " - " + errorResponse);
            }
            
            // 读取响应
            return readInputStream(connection.getInputStream());
            
        } finally {
            connection.disconnect();
        }
    }
    
    /**
     * 发送GET请求
     */
    private String sendGetRequest(String urlString) throws IOException {
        URL url = new URL(urlString);
        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
        
        try {
            // 设置请求方法和属性
            connection.setRequestMethod("GET");
            connection.setRequestProperty("Accept", "application/json");
            connection.setConnectTimeout(connectTimeout);
            connection.setReadTimeout(readTimeout);
            
            // 检查响应状态码
            int responseCode = connection.getResponseCode();
            if (responseCode != HttpURLConnection.HTTP_OK) {
                String errorResponse = readErrorStream(connection);
                throw new IOException("HTTP Error: " + responseCode + " - " + errorResponse);
            }
            
            // 读取响应
            return readInputStream(connection.getInputStream());
            
        } finally {
            connection.disconnect();
        }
    }
    
    /**
     * 读取输入流
     */
    private String readInputStream(InputStream inputStream) throws IOException {
        StringBuilder response = new StringBuilder();
        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
            String line;
            while ((line = reader.readLine()) != null) {
                response.append(line);
            }
        }
        return response.toString();
    }
    
    /**
     * 读取错误流
     */
    private String readErrorStream(HttpURLConnection connection) {
        try {
            InputStream errorStream = connection.getErrorStream();
            if (errorStream != null) {
                return readInputStream(errorStream);
            }
        } catch (IOException e) {
            // 忽略读取错误流时的异常
        }
        return "Unknown error";
    }
    
    /**
     * 构建JSON请求体
     */
    private String buildJsonRequest(OllamaRequest request) {
        StringBuilder json = new StringBuilder();
        json.append("{");
        json.append("\"model\":\"").append(escapeJson(request.getModel())).append("\",");
        json.append("\"prompt\":\"").append(escapeJson(request.getPrompt())).append("\"");
        // 显式指定是否流式，避免默认开启流模式导致解析只取到第一块
        json.append(",\"stream\":").append(request.isStream() ? "true" : "false");
        json.append("}");
        return json.toString();
    }
    
    /**
     * 构建聊天JSON请求体
     */
    private String buildChatJsonRequest(OllamaChatRequest request) {
        StringBuilder json = new StringBuilder();
        json.append("{");
        json.append("\"model\":\"").append(escapeJson(request.getModel())).append("\",");
        json.append("\"messages\":[");
        for (int i = 0; i < request.getMessages().size(); i++) {
            if (i > 0) json.append(",");
            OllamaChatMessage msg = request.getMessages().get(i);
            json.append("{\"role\":\"").append(escapeJson(msg.getRole())).append("\",");
            json.append("\"content\":\"").append(escapeJson(msg.getContent())).append("\"}");
        }
        json.append("]");
        if (request.isStream()) {
            json.append(",\"stream\":true");
        }
        json.append("}");
        return json.toString();
    }
    
    /**
     * 简单的JSON转义
     */
    private String escapeJson(String str) {
        if (str == null) return "";
        return str.replace("\\", "\\\\")
                  .replace("\"", "\\\"")
                  .replace("\n", "\\n")
                  .replace("\r", "\\r")
                  .replace("\t", "\\t");
    }
    
    /**
     * 解析响应
     */
    private OllamaResponse parseResponse(String jsonResponse) {
        String response = extractJsonValue(jsonResponse, "response");
        boolean done = "true".equals(extractJsonValue(jsonResponse, "done"));
        return new OllamaResponse(response, done);
    }
    
    /**
     * 解析聊天响应
     */
    private OllamaChatResponse parseChatResponse(String jsonResponse) {
        // 提取message对象中的content
        String messageStart = "\"message\":{";
        int messageIndex = jsonResponse.indexOf(messageStart);
        if (messageIndex != -1) {
            String messageJson = jsonResponse.substring(messageIndex + messageStart.length());
            int messageEnd = findMatchingBrace(messageJson);
            if (messageEnd != -1) {
                String messageContent = messageJson.substring(0, messageEnd);
                String content = extractJsonValue(messageContent, "content");
                String role = extractJsonValue(messageContent, "role");
                boolean done = "true".equals(extractJsonValue(jsonResponse, "done"));
                return new OllamaChatResponse(new OllamaChatMessage(role, content), done);
            }
        }
        return new OllamaChatResponse(new OllamaChatMessage("assistant", ""), true);
    }
    
    /**
     * 改进的JSON值提取
     */
    private String extractJsonValue(String json, String key) {
        String searchKey = "\"" + key + "\":\"";
        int startIndex = json.indexOf(searchKey);
        if (startIndex == -1) {
            // 尝试布尔值或数字
            searchKey = "\"" + key + "\":";
            startIndex = json.indexOf(searchKey);
            if (startIndex == -1) return null;
            startIndex += searchKey.length();
            int endIndex = json.indexOf(",", startIndex);
            if (endIndex == -1) endIndex = json.indexOf("}", startIndex);
            if (endIndex == -1) return null;
            String value = json.substring(startIndex, endIndex).trim();
            // 移除可能的引号
            if (value.startsWith("\"") && value.endsWith("\"")) {
                value = value.substring(1, value.length() - 1);
            }
            return value;
        }

        startIndex += searchKey.length();
        int endIndex = startIndex;
        boolean inString = true; // 我们已经在字符串内了

        while (endIndex < json.length()) {
            char c = json.charAt(endIndex);

            if (c == '\\' && inString) {
                endIndex += 2; // 跳过转义字符
                continue;
            }

            if (c == '"' && inString) {
                // 找到字符串结束
                break;
            }

            endIndex++;
        }

        if (endIndex >= json.length()) return null;
        String result = json.substring(startIndex, endIndex);

        // 处理Unicode转义序列
        result = unescapeJson(result);

        return result;
    }

    /**
     * 处理JSON转义序列
     */
    private String unescapeJson(String str) {
        if (str == null) return null;
        return str.replace("\\\"", "\"")
                  .replace("\\\\", "\\")
                  .replace("\\n", "\n")
                  .replace("\\r", "\r")
                  .replace("\\t", "\t")
                  .replace("\\u003c", "<")
                  .replace("\\u003e", ">")
                  .replace("\\u0026", "&");
    }
    
    /**
     * 找到匹配的大括号
     */
    private int findMatchingBrace(String json) {
        int braceCount = 1;
        for (int i = 0; i < json.length(); i++) {
            char c = json.charAt(i);
            if (c == '{') braceCount++;
            else if (c == '}') {
                braceCount--;
                if (braceCount == 0) return i;
            }
        }
        return -1;
    }
    
    /**
     * 关闭客户端，释放资源
     */
    public void close() {
        if (executorService != null && !executorService.isShutdown()) {
            executorService.shutdown();
        }
    }
}
